package com.peanut.sql.lineage.postgresql;

import com.alibaba.druid.DbType;
import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLName;
import com.alibaba.druid.sql.ast.SQLObject;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.expr.*;
import com.alibaba.druid.sql.ast.statement.*;
import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGInsertStatement;
import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectQueryBlock;
import com.alibaba.druid.sql.parser.ParserException;
import com.alibaba.druid.sql.visitor.SchemaStatVisitor;
import com.alibaba.druid.util.JdbcConstants;
import com.peanut.utils.Neo4jUtil;
import org.apache.commons.lang3.StringUtils;
import org.neo4j.driver.Result;
import org.neo4j.driver.Session;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.*;

import static org.neo4j.driver.Values.parameters;

/****
 author:yifan
 date:2016
 description:
 **/
public class PGSQLLineage {

    /**
     * 适用于目标表只有1个
     * insert into {tableName} select;create table {tableName} as select
     * todo 强化对with语法的解析,部分函数的解析
     *
     * @param sqlStr
     * @return
     * @throws Exception
     */
    public LinkedHashMap<Column, Set<Vertex>> getColumnLineage(String sqlStr) throws Exception {
        LinkedHashMap<Column, Set<Vertex>> result = new LinkedHashMap<>();
        HashMap<String, Column[]> metaData = getMetaData();
        DbType dbType = JdbcConstants.POSTGRESQL;
        SQLStatement sqlStatement = SQLUtils.parseSingleStatement(sqlStr, dbType);
        SchemaStatVisitor statVisitor = SQLUtils.createSchemaStatVisitor(dbType);
        sqlStatement.accept(statVisitor);
        if (sqlStatement instanceof PGInsertStatement || sqlStatement instanceof SQLCreateTableStatement) {
            SQLSelectQuery createTableStatementQuery = null;
            SQLSelect sqlSelect ;
            SQLExpr exprTarget ;
            if (sqlStatement instanceof PGInsertStatement) {
                sqlSelect = ((PGInsertStatement) sqlStatement).getQuery();
                exprTarget = ((PGInsertStatement) sqlStatement).getTableSource().getExpr();
            } else {
                sqlSelect = ((SQLCreateTableStatement) sqlStatement).getSelect();
                createTableStatementQuery = sqlSelect.getQuery();
                exprTarget = ((SQLCreateTableStatement) sqlStatement).getTableSource().getExpr();
            }
            if (exprTarget instanceof SQLPropertyExpr) {
                String schemaName = ((SQLPropertyExpr) exprTarget).getOwnerName();
                String tableName = ((SQLPropertyExpr) exprTarget).getName();
                String fullTableName = "postgres." + schemaName + "." + tableName;
                //        将表结构初始化到metadata中
                if (createTableStatementQuery != null) {
                    updateMetadata(createTableStatementQuery, fullTableName, metaData);
                }
                Column[] targetColumns = metaData.get(fullTableName);
                List<Set<Vertex>> index = new ArrayList<>();
                index = getIndex(sqlSelect.getQuery(), index);
                if (targetColumns.length != index.size()) {
                    throw new ParserException("selected column number not equal to target table!");
                }
                for (int i = 0; i < targetColumns.length; i++) {
                    result.put(targetColumns[i], index.get(i));
                }
            } else {
                throw new ParserException("can not find target table properly!");
            }
        } else {
            throw new ParserException("no lineage need to analyze!");
        }
//        图形化展示
        toGraph(result);
        return result;
    }

    private void updateMetadata(SQLSelectQuery query, String fullTableName, HashMap<String, Column[]> metaData) {
        List<SQLSelectItem> selectList = null;
        if (query instanceof SQLSelectQueryBlock) {
            selectList = ((SQLSelectQueryBlock) query).getSelectList();
        } else if (query instanceof SQLUnionQuery) {
            SQLSelectQuery sqlSelectQuery = ((SQLUnionQuery) query).getRelations().get(0);
            updateMetadata(sqlSelectQuery, fullTableName, metaData);
        }
        if (selectList != null) {
            Column[] columns = new Column[selectList.size()];
            String[] preColumnDefine = StringUtils.split(fullTableName, '.');
            for (int i = 0; i < selectList.size(); i++) {
                SQLSelectItem sqlSelectItem = selectList.get(i);
                String columnName = sqlSelectItem.getExpr().toString();
                String columnAlias = sqlSelectItem.getAlias();
                String columnNameFinal = columnAlias != null ? columnAlias : columnName;
                columns[i] = new Column(columnNameFinal, null, null, preColumnDefine[2], preColumnDefine[1], preColumnDefine[0]);
            }
            metaData.put(fullTableName, columns);
        }
    }

    public List<Set<Vertex>> getIndex(SQLSelectQuery sqlSelectQuery, List<Set<Vertex>> index) throws Exception {
        if (sqlSelectQuery instanceof SQLSelectQueryBlock) {
            List<SQLSelectItem> selectList = ((SQLSelectQueryBlock) sqlSelectQuery).getSelectList();
            for (SQLSelectItem sqlSelectItem : selectList) {
                SQLExpr selectItemExpr = sqlSelectItem.getExpr();
                LinkedList<Vertex> sourceFields = new LinkedList<>();
                getSourceFieldsByItemExpr(selectItemExpr, sourceFields);
                HashSet<Vertex> sourceFieldsSet = new HashSet<>();
//                    当源字段全部为常量时，来源仅显示常量标识，其他的去除常量显示
                boolean isAllConstant = true;
                for (Vertex sourceField : sourceFields) {
                    if (!(sourceField instanceof Constant)) {
                        isAllConstant = false;
                        sourceFieldsSet.add(sourceField);
                    }
                }
                if (isAllConstant) {
                    sourceFieldsSet.clear();
//                        目前不展示常量的值，仅作为标识
                    sourceFieldsSet.add(new Constant("CONSTANT"));
                }
                index.add(sourceFieldsSet);
            }
        } else if (sqlSelectQuery instanceof SQLUnionQueryTableSource || sqlSelectQuery instanceof SQLUnionQuery) {
            SQLUnionQuery sqlUnionQuery = null;
            if (sqlSelectQuery instanceof SQLUnionQueryTableSource) {
                sqlUnionQuery = ((SQLUnionQueryTableSource) sqlSelectQuery).getUnion();
            }
            if (sqlSelectQuery instanceof SQLUnionQuery) {
                sqlUnionQuery = (SQLUnionQuery) sqlSelectQuery;
            }
            List<SQLSelectQuery> sqlUnionQueryRelations = sqlUnionQuery.getRelations();
            for (SQLSelectQuery sqlUnionQueryRelation : sqlUnionQueryRelations) {
                List<Set<Vertex>> result = new ArrayList<>();
                getIndex(sqlUnionQueryRelation, result);
                if (index.size() != 0 && result.size() != 0) {
                    List<Set<Vertex>> resultMerge = new ArrayList<>();
                    for (int i = 0; i < result.size(); i++) {
                        HashSet<Vertex> sourceVertex = new HashSet<>();
                        sourceVertex.addAll(result.get(i));
                        sourceVertex.addAll(index.get(i));
                        resultMerge.add(sourceVertex);
                    }
                    index = resultMerge;

                } else {
                    if (result.size() != 0) {
                        index = result;
                    }
                }
            }
        } else {
            throw new ParserException("unSupport sql!");
        }
        return index;
    }

    private LinkedList<Vertex> getSourceFieldsByItemExpr(SQLExpr expr, LinkedList<Vertex> sourceFields) throws
            Exception {
        if (expr instanceof SQLPropertyExpr || expr instanceof SQLIdentifierExpr) {
            String fieldName = null;
            SQLObject resolvedOwnerObject = null;
            if (expr instanceof SQLPropertyExpr) {
                fieldName = ((SQLPropertyExpr) expr).getName();
                resolvedOwnerObject = ((SQLPropertyExpr) expr).getResolvedOwnerObject();
            }
            if (expr instanceof SQLIdentifierExpr) {
                fieldName = ((SQLIdentifierExpr) expr).getName();
                resolvedOwnerObject = ((SQLIdentifierExpr) expr).getResolvedOwnerObject();
            }

            if (resolvedOwnerObject instanceof SQLExprTableSource) {
                SQLExpr propertyExpr = ((SQLExprTableSource) resolvedOwnerObject).getExpr();
                if (propertyExpr instanceof SQLPropertyExpr) {
                    Column sourceColumn = new Column(fieldName, null, null, ((SQLPropertyExpr) propertyExpr).getName(), ((SQLPropertyExpr) propertyExpr).getOwnerName(), "postgres");
                    sourceFields.add(sourceColumn);
                }
            } else if (resolvedOwnerObject instanceof SQLSubqueryTableSource) {
                SQLSelect sqlSelect = ((SQLSubqueryTableSource) resolvedOwnerObject).getSelect();
                SQLSelectQueryBlock queryBlock = sqlSelect.getQueryBlock();
                SQLSelectItem selectItem = queryBlock.findSelectItem(fieldName);
                if (selectItem == null) {
                    throw new Exception("SQL error! not found column");
                }
                getSourceFieldsByItemExpr(selectItem.getExpr(), sourceFields);
            } else if (resolvedOwnerObject instanceof SQLUnionQueryTableSource) {
                SQLUnionQuery sqlUnionQuery = ((SQLUnionQueryTableSource) resolvedOwnerObject).getUnion();
                List<SQLSelectQuery> relations = sqlUnionQuery.getRelations();
                PGSelectQueryBlock queryBlock = (PGSelectQueryBlock) relations.get(0);
                SQLSelectItem selectItem = queryBlock.findSelectItem(fieldName);
                List<SQLSelectItem> selectList = queryBlock.getSelectList();
                int position = selectList.indexOf(selectItem);
                for (SQLSelectQuery relation : relations) {
                    if (relation instanceof PGSelectQueryBlock) {
                        List<SQLSelectItem> sqlSelectItems = ((PGSelectQueryBlock) relation).getSelectList();
                        SQLSelectItem sqlSelectItem = sqlSelectItems.get(position);
                        getSourceFieldsByItemExpr(sqlSelectItem.getExpr(), sourceFields);
                    }

                }
            } else if (resolvedOwnerObject instanceof SQLWithSubqueryClause.Entry) {
                List<SQLName> columns = ((SQLWithSubqueryClause.Entry) resolvedOwnerObject).getColumns();
//  当前只处理columns有值的情况
                if (columns.size() != 0) {
                    int position = -1;
                    for (int i = 0; i < columns.size(); i++) {
                        if (columns.get(i).getSimpleName().equals(fieldName)) {
                            position = i;
                            break;
                        }
                    }
                    if (position == -1) {
                        throw new Exception(String.format("not found column %s!", fieldName));
                    }
                    SQLSelectQuery query = ((SQLWithSubqueryClause.Entry) resolvedOwnerObject).getSubQuery().getQuery();
                    List<Set<Vertex>> index = new ArrayList<>();
                    List<Set<Vertex>> indexResult = getIndex(query, index);
                    if (indexResult.size() > position) {
                        Set<Vertex> vertices = indexResult.get(position);
                        sourceFields.addAll(vertices);
                    }
                }
            } else {
                throw new Exception("unSupport resolvedOwnerObject type!");
            }

        } else if (expr instanceof SQLMethodInvokeExpr || expr instanceof SQLAggregateExpr) {
            List<SQLExpr> arguments = ((SQLMethodInvokeExpr) expr).getArguments();
            if (arguments.size() == 0) {
                sourceFields.add(new PGFunction(((SQLMethodInvokeExpr) expr).getMethodName(), null, null, null, null));
            } else {
                for (SQLExpr argument : arguments) {
                    getSourceFieldsByItemExpr(argument, sourceFields);
                }
            }
        } else if (expr instanceof SQLArrayExpr) {
            List<SQLExpr> values = ((SQLArrayExpr) expr).getValues();
            for (SQLExpr value : values) {
                getSourceFieldsByItemExpr(value, sourceFields);
            }
        } else if (expr instanceof SQLCharExpr) {
            String text = ((SQLCharExpr) expr).getText();
            sourceFields.add(new Constant(text));
        } else if (expr instanceof SQLIntegerExpr) {
            Number number = ((SQLIntegerExpr) expr).getNumber();
            sourceFields.add(new Constant(number.toString()));
        } else if (expr instanceof SQLCaseExpr) {
            List<SQLCaseExpr.Item> items = ((SQLCaseExpr) expr).getItems();
            for (SQLCaseExpr.Item item : items) {
                SQLExpr conditionExpr = item.getConditionExpr();
                getSourceFieldsByItemExpr(conditionExpr, sourceFields);
                SQLExpr valueExpr = item.getValueExpr();
                getSourceFieldsByItemExpr(valueExpr, sourceFields);
            }
            SQLExpr elseExpr = ((SQLCaseExpr) expr).getElseExpr();
            getSourceFieldsByItemExpr(elseExpr, sourceFields);

        } else if (expr instanceof SQLBinaryOpExpr) {
            SQLExpr left = ((SQLBinaryOpExpr) expr).getLeft();
            getSourceFieldsByItemExpr(left, sourceFields);
            SQLExpr right = ((SQLBinaryOpExpr) expr).getRight();
            getSourceFieldsByItemExpr(right, sourceFields);
        }

        return sourceFields;
    }

    private void toGraph(LinkedHashMap<Column, Set<Vertex>> result) {
        Session neo4jSession = Neo4jUtil.getNeo4jSession();
        neo4jSession.run("MATCH (n) DETACH DELETE n");
        Set<Map.Entry<Column, Set<Vertex>>> entries = result.entrySet();
        for (Map.Entry<Column, Set<Vertex>> entry : entries) {
            Column targetColumn = entry.getKey();
            Set<Vertex> sourceVertex = entry.getValue();
            createTableVertex(targetColumn.getSchemaName(), targetColumn.getTableName());
            createFieldVertex(targetColumn.getSchemaName(), targetColumn.getTableName(), targetColumn.getColumnName());
            for (Vertex vertex : sourceVertex) {
                if (vertex instanceof Column) {
                    createTableVertex(((Column) vertex).getSchemaName(), ((Column) vertex).getTableName());
                    createFieldVertex(((Column) vertex).getSchemaName(), ((Column) vertex).getTableName(), ((Column) vertex).getColumnName());
                    neo4jSession.run("MATCH (a:FieldVertex),(b:FieldVertex)" +
                                    "WHERE a.dbName = $srcDbName " +
                                    "and a.tableName = $srcTableName " +
                                    "and a.fieldName= $srcFieldName " +
                                    "and b.dbName = $tgDbName " +
                                    "and b.tableName = $tgTableName " +
                                    "and b.fieldName= $tgFieldName " +
                                    "CREATE (a)-[r:INPUT]->(b)",
                            parameters("srcDbName", ((Column) vertex).getSchemaName(),
                                    "srcTableName", ((Column) vertex).getTableName(), "srcFieldName", ((Column) vertex).getColumnName()
                                    , "tgDbName", targetColumn.getSchemaName(),
                                    "tgTableName", targetColumn.getTableName(), "tgFieldName",
                                    targetColumn.getColumnName()));
                } else if (vertex instanceof Constant) {
                    neo4jSession.run("create (n:Constant{constantValue:$constantValue,tgDbName:$tgDbName,tgTableName:$tgTableName,tgFieldName:$tgFieldName}) ",
                            parameters("constantValue", ((Constant) vertex).getConstantValue()
                                    , "tgDbName", targetColumn.getSchemaName(), "tgTableName", targetColumn.getTableName(), "tgFieldName", targetColumn.getColumnName()));
                    neo4jSession.run("MATCH (a:FieldVertex),(b:Constant)" +
                                    "WHERE a.dbName = $srcDbName " +
                                    "and a.tableName = $srcTableName " +
                                    "and a.fieldName= $srcFieldName " +
                                    "and b.tgDbName = $tgDbName " +
                                    "and b.tgTableName = $tgTableName " +
                                    "and b.tgFieldName = $tgFieldName " +
                                    "CREATE (b)-[r:INPUT]->(a)",
                            parameters("srcDbName", targetColumn.getSchemaName(),
                                    "srcTableName", targetColumn.getTableName(), "srcFieldName", targetColumn.getColumnName()
                                    , "tgDbName", targetColumn.getSchemaName(),
                                    "tgTableName", targetColumn.getTableName(), "tgFieldName",
                                    targetColumn.getColumnName()));
                } else {
                    neo4jSession.run("create (n:PGFunction{functionName:$functionName,tgDbName:$tgDbName,tgTableName:$tgTableName,tgFieldName:$tgFieldName}) ",
                            parameters("functionName", ((PGFunction) vertex).getFunctionName()
                                    , "tgDbName", targetColumn.getSchemaName(), "tgTableName", targetColumn.getTableName(), "tgFieldName", targetColumn.getColumnName()));
                    neo4jSession.run("MATCH (a:FieldVertex),(b:PGFunction)" +
                                    "WHERE a.dbName = $srcDbName " +
                                    "and a.tableName = $srcTableName " +
                                    "and a.fieldName= $srcFieldName " +
                                    "and b.tgDbName = $tgDbName " +
                                    "and b.tgTableName = $tgTableName " +
                                    "and b.tgFieldName= $tgFieldName " +
                                    "CREATE (b)-[r:INPUT]->(a)",
                            parameters("srcDbName", targetColumn.getSchemaName(),
                                    "srcTableName", targetColumn.getTableName(), "srcFieldName", targetColumn.getColumnName()
                                    , "tgDbName", targetColumn.getSchemaName(),
                                    "tgTableName", targetColumn.getTableName(), "tgFieldName",
                                    targetColumn.getColumnName()));
                }
            }
        }

        neo4jSession.run("MATCH (a:TableVertex),(b:FieldVertex)" +
                "WHERE a.dbName = b.dbName " +
                "and a.tableName = b.tableName " +
                "CREATE (a)-[r:OWNER]->(b)");
        neo4jSession.close();
    }

    private void createFieldVertex(String dbName, String tableName, String fieldName) {
        String querySql = String.format("match (a:FieldVertex) where a.dbName='%s' and a.tableName='%s' and a.fieldName='%s' return a", dbName, tableName, fieldName);
        if (!isNodeExists(querySql)) {
            Neo4jUtil.getNeo4jSession().run("create (n:FieldVertex{dbName:$dbName,tableName:$tableName,fieldName:$fieldName}) ",
                    parameters("dbName", dbName,
                            "tableName", tableName, "fieldName", fieldName));
        }
    }

    private void createTableVertex(String dbName, String tableName) {
        String querySql = String.format("match (a:TableVertex) where a.dbName='%s' and a.tableName='%s' return a", dbName, tableName);
        if (!isNodeExists(querySql)) {
            Neo4jUtil.getNeo4jSession().run("create (n:TableVertex{dbName:$dbName,tableName:$tableName}) ",
                    parameters("dbName", dbName,
                            "tableName", tableName));
        }
    }

    private boolean isNodeExists(String sql) {
        Result resultExists = Neo4jUtil.getNeo4jSession().run(sql);
        return resultExists.list().size() != 0 ;
    }

    private HashMap<String, Column[]> getMetaData() {
        HashMap<String, Column[]> metaData = new HashMap<>();
        try {
            Class.forName("org.postgresql.Driver");
            Connection conn = DriverManager.getConnection("jdbc:postgresql://localhost:5432/postgres", "postgres", "123456");
            Statement  stmt = conn.createStatement();
            ResultSet rs = stmt.executeQuery("select table_catalog,table_schema,table_name,string_agg(column_name,',') as columns from information_schema.columns group by table_catalog,table_schema,table_name");
            while (rs.next()) {
                String tableCatalog = rs.getString("table_catalog");
                String tableSchema = rs.getString("table_schema");
                String tableName = rs.getString("table_name");
                String[] columns = rs.getString("columns").split(",");
                Column[] columnObjects = new Column[columns.length];
                for (int i = 0; i < columns.length; i++) {
                    columnObjects[i] = new Column(columns[i], null, null, tableName, tableSchema, tableCatalog);
                }
                metaData.put(tableCatalog + "." + tableSchema + "." + tableName, columnObjects);
            }
            rs.close();
            stmt.close();
            conn.close();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return metaData;
    }
}
